import base64
import zipfile
import psutil
import numpy as np
import torch
import matplotlib.pyplot as plt
import cv2
from routes.GlobalState import global_state
from segment_anything.utils.transforms import ResizeLongestSide
import sys
import os
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator, SamPredictor
from service.segment.AutomaticDemonstration import IncomingCoordinatePoint, update_mask, init_model, save_masked_image, \
    create_class_index_mapping, mouse_click_back, save_and_add_picture_yolo, save_and_add_picture_xml, \
    save_and_add_picture_coco, generate_masked_image, save_and_add_picture_voc
from service.segment.runSeg import show_mask,show_points,show_box,show_anns,generate_image_number,get_resize_transform
# from service.segment.runSeg import Simple_coordinate_new_image, Simple_coordinate_new_image2,Simple_coordinate_new_image3
# from service.segment.runSeg import load_image, load_sam_model, process_image, predict_mask,predict_mask2,prepare_image
# from service.segment.runSeg import get_batched_input,batch_process_images,generate_masks
from routes.app import app
from DBService.User_Service import UserService
from flask import request
from DBService.Working_Service import WorkingService
from flask import request, send_file
from PIL import Image
from flask import g
import redis
import pickle
import io
from flask import jsonify
from werkzeug.utils import secure_filename


@app.route('/segment/init_model', methods=['POST'])
def init_model_route():
    # 从请求中获取参数
    image_name = request.args.get('image_name')

    # 初始化模型
    # 初始化模型
    global_state.image, global_state.sam, global_state.predictor, global_state.input_point, \
        global_state.input_label, global_state.input_box_arr, global_state.mask_input, \
        global_state.multimask_output, global_state.input_stop,global_state. image_path= init_model(
        image_name)
    print("模型初始化完成")

    # 返回初始化成功的消息
    return jsonify({"message": "模型初始化成功"})

#撤销最后添加的操作
@app.route('/segment/undo_last_operation', methods=['POST'])
def undo_last_operation():
    # 调用mouse_click_back函数
    mouse_click_back()
    return jsonify({"message": "最后一次操作已撤销"})

#持续获取前端传来的坐标数据
@app.route('/segment/update_point', methods=['POST'])
def update_point_route():
    # 从全局变量中获取模型和其他参数
    global global_state
    # 从请求中获取参数
    x = int(request.json.get('x'))
    y = int(request.json.get('y'))
    label = int(request.json.get('label'))

    # 更新点和标签
    IncomingCoordinatePoint(x, y, label)
    print("点和标签更新成功")

    return jsonify({"message": "点和标签更新成功"})

@app.route('/segment/update_mask', methods=['POST'])
def update_mask_route():
    # 从全局变量中获取模型和其他参数
    global global_state

    # 更新mask
    global_state.masks, global_state.scores, global_state.logits, global_state.mask_input, global_state.input_stop, global_state.input_point_arr, global_state.input_label_arr = update_mask(global_state.image, global_state.sam, global_state.predictor, global_state.input_point, global_state.input_label, global_state.mask_input, global_state.multimask_output, global_state.input_stop, global_state.input_box_arr)
    print("mask更新成功")
    # 绘制图像
    fig, ax = plt.subplots(figsize=(10, 10))
    ax.imshow(global_state.image)
    show_mask(global_state.masks, ax)
    show_points(global_state.input_point_arr, global_state.input_label_arr, ax)
    #绘制标记框（如果有）
    show_box(global_state.input_box_arr, ax)
    plt.axis('off')

    # 移除图像周围的空白边缘
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    # 获取global_state.image的分辨率
    dpi = global_state.image.shape[1] / 10  # 假设figsize的宽度为10
    # 将图像保存为一个临时文件
    img_io = io.BytesIO()
    plt.savefig(img_io, format='png', bbox_inches='tight', pad_inches=0,dpi=dpi)
    img_io.seek(0)
    # 将文件转换为Base64格式
    img_base64 = base64.b64encode(img_io.getvalue()).decode('utf-8')
    plt.close(fig)
    # 将Base64格式的图片发送给前端
    return jsonify({"image": img_base64})
#保存可行图片
@app.route('/segment/input_box_arr', methods=['POST'])
def update_input_box_arr():
    global global_state
    # 从请求的JSON数据中获取参数
    data = request.get_json()
    print(data)
    x1 = int(data.get('x1'))
    y1 = int(data.get('y1'))
    x2 = int(data.get('x2'))
    y2 = int(data.get('y2'))
    # 更新box
    global_state.input_box_arr = np.array([x1, y1, x2, y2])
    print("box更新成功"+str(global_state.input_box_arr))

    return jsonify({"message": "box更新成功"})
#更改g.input_stop的值
@app.route('/segment/input_stop')
def input_stop():
    global global_state
    global_state.input_stop=True
    return "success"
# @app.route('/segment/save_masked_image')
# def save_masked_image_route():
#     # 从全局变量g中获取模型和其他参数
#     image = g.image
#     masks=g.masks
#
#     # 保存图片
#     save_path1 = save_masked_image(image, masks)
#     # 将文件路径发送给前端
#     return save_path1

#获取用户所有图片
@app.route('/segment/get_working')
def get_working_all():
    uid = request.args.get('uid')
    workings = WorkingService().get_working(uid)
    return workings


@app.route('/memory')
def get_memory_usage():
    memory_percent = psutil.virtual_memory().percent
    return {'memory_percent': memory_percent}

#接受上传图片并将其保存到service/segment/image文件夹
@app.route('/segment/upload', methods=['POST'])
def upload_file():
    print("Files:", request.files)
    print("Form data:", request.form)
    file = request.files.get('File')
    if file:
        filename = secure_filename(file.filename)
        # 检查文件名是否为空
        if file.filename == '':
            return jsonify({"error": "没有选择文件"}), 402

        # 检查文件类型
        if not allowed_file(file.filename):
            return jsonify({"error": "不允许的文件类型"}), 403
        # 保存文件
        dir_path = os.path.abspath(os.path.join('..', 'ThinkStation','service', 'segment', 'image'))
        print(dir_path)
        if not os.path.exists(dir_path):
            os.makedirs(dir_path)
        save_path = os.path.join(dir_path, filename)
        try:
            file.save(save_path)
            print('图片已保存' + save_path)
            file_size = os.path.getsize(save_path)
            uid = request.form.get('uid')
            # 将图片信息存入数据库
            relative_save_path = os.path.relpath(save_path, start=os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
            image=WorkingService().add_DB_image(uid, relative_save_path,filename)
            if image is not None:
                return jsonify({"message": "文件上传成功", "filename": filename, "size": file_size}), 200  # 使用 jsonify 返回 JSON 响应
            else:
                return jsonify({"error": "文件上传失败"}), 500
        except Exception as e:

            return jsonify({"error": "保存文件或存入数据库时发生错误: " + str(e)}), 500
    else:
        return jsonify({"message": "请求中没有文件部分."}), 400


#根据用户的uid获取用户所有的图片-image中
@app.route('/segment/get_image')
def get_image():
    uid = request.args.get('uid')
    images_path = WorkingService().get_DB_image(uid)
    return images_path

##删除图片-image中
@app.route('/segment/delete_image')
def delete_image():
    id = request.args.get('id')
    image = WorkingService().del_DB_image(id)
    print("图片已删除:" + image.filename)
    return jsonify({"message": "文件删除成功"}), 200  # 使用 jsonify 返回 JSON 响应



#根据用户uid来查询所有制作好的数据-Working中
@app.route('/segment/get_working')
def get_working_uid():
    uid = request.args.get('uid')
    workings = WorkingService().get_working(uid)
    return workings
#根据前端传来的图片id序列来在数据库中查询需要导出yolo格式的的workings序列

#图片满足条件后保存工作图片到指定文件夹并将图片信息存入数据库
@app.route('/segment/save_and_add_picture', methods=['POST'])
def save_and_add_picture():
    # 从全局变量global_state中获取模型和其他参数
    image = global_state.image
    masks = global_state.masks
    uid = request.json.get('id')
    print("uid:"+uid)
    classification_old = request.json.get('classification')
    class_index_mapping = create_class_index_mapping()
    classification = class_index_mapping.get(classification_old)
    #保存原始图片
    save_path_old =global_state.image_path
    print("save_path_old"+save_path_old)
    # 生成模框图片
    img_io, xmin, ymin, xmax, ymax = generate_masked_image(image, masks)
    # 保存模框图片
    save_path_new = save_masked_image(img_io)
    print(save_path_new)
    # 计算图片大小
    size = os.path.getsize(save_path_new)

    # 将图片信息存入数据库
    WorkingService().add_working(uid, size, save_path_old, xmin, ymin, xmax, ymax,classification,save_path_new)
    # 将模框路径和原始图片路径发送给前端
    return jsonify({"message": "图片保存成功", "pathNew": save_path_old, "pathOld": save_path_new}), 200
#右下角展示
@app.route('/segment/show_anns', methods=['POST'])
def show_anns_route():
    # 从全局变量global_state中获取模型和其他参数
    image = global_state.image
    masks = global_state.masks
    img_io, xmin, ymin, xmax, ymax = generate_masked_image(image, masks)
    #将img_io转换为base64格式
    img_base64 = base64.b64encode(img_io.getvalue()).decode('utf-8')
    #发送给前端
    return jsonify({"image": img_base64})
# @app.route('/segment/get_working_toFile', methods=['POST'])
# def get_working_toFile():
#     id_list = request.json.get('id_list')
#     workings_toFile = WorkingService().get_working_toFile(id_list)
#     return workings_toFile
#导出yolo格式的数据
@app.route('/segment/export_yolo', methods=['POST'])
def export_yolo():
    id_list = request.json.get('id_list')
    print(id_list)
    workings_toFile = WorkingService().get_working_toFile(id_list)
    #调用save_and_add_picture_yolo函数
    #print("workings_toFile:"+workings_toFile)
    yolo_path,image_dir=save_and_add_picture_yolo(workings_toFile)
    print("yolo_path:"+yolo_path)
    print("image_dir:"+image_dir)

    zip_path = image_dir + '.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(image_dir):
            for file in files:
                zipf.write(os.path.join(root, file),
                           os.path.relpath(os.path.join(root, file),
                                           os.path.join(image_dir, '..')))

    # 将zip文件发送到客户端
    response = send_file(zip_path, mimetype='application/zip', as_attachment=True)
    response.headers["Content-Disposition"] = "attachment; filename={}".format(os.path.basename(zip_path))
    return response
#导出xml格式的数据
@app.route('/segment/export_xml', methods=['POST'])
def export_xml():
    id_list = request.json.get('id_list')
    print(id_list)
    workings_toFile = WorkingService().get_working_toFile(id_list)
    #调用save_and_add_picture_yolo函数
    xml_path,image_dir=save_and_add_picture_xml(workings_toFile)
    print("xml_path:"+xml_path)
    print("image_dir:"+image_dir)
    zip_path = image_dir + '.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(image_dir):
            for file in files:
                zipf.write(os.path.join(root, file),
                           os.path.relpath(os.path.join(root, file),
                                           os.path.join(image_dir, '..')))

    # 将zip文件发送到客户端
    response = send_file(zip_path, mimetype='application/zip', as_attachment=True)
    response.headers["Content-Disposition"] = "attachment; filename={}".format(os.path.basename(zip_path))
    return response
#导出coco格式的数据
@app.route('/segment/export_coco', methods=['POST'])
def export_coco():
    id_list = request.json.get('id_list')
    print(id_list)
    workings_toFile = WorkingService().get_working_toFile(id_list)
    #调用save_and_add_picture_yolo函数
    coco_path,image_dir=save_and_add_picture_coco(workings_toFile)
    print("coco_path:"+coco_path)
    print("image_dir:"+image_dir)
    #根据image_dir的路径来压缩这个文件包并将这个文件包发送给前端
    #从该目录创建一个zip文件
    zip_path = image_dir + '.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(image_dir):
            for file in files:
                zipf.write(os.path.join(root, file),
                           os.path.relpath(os.path.join(root, file),
                                           os.path.join(image_dir, '..')))

    # 将zip文件发送到客户端
    response = send_file(zip_path, mimetype='application/zip', as_attachment=True)
    response.headers["Content-Disposition"] = "attachment; filename={}".format(os.path.basename(zip_path))
    return response
#导出voct格式的数据
@app.route('/segment/export_voc', methods=['POST'])
def export_voc():
    id_list = request.json.get('id_list')
    print(id_list)
    workings_toFile = WorkingService().get_working_toFile(id_list)
    #调用save_and_add_picture_voc函数
    voc_path,image_dir=save_and_add_picture_voc(workings_toFile)
    print("voc_path:"+voc_path)
    print("image_dir:"+image_dir)
    #根据image_dir的路径来压缩这个文件包并将这个文件包发送给前端
    #从该目录创建一个zip文件
    zip_path = image_dir + '.zip'
    with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
        for root, dirs, files in os.walk(image_dir):
            for file in files:
                zipf.write(os.path.join(root, file),
                           os.path.relpath(os.path.join(root, file),
                                           os.path.join(image_dir, '..')))
    # 将zip文件发送到客户端
    response = send_file(zip_path, mimetype='application/zip', as_attachment=True)
    response.headers["Content-Disposition"] = "attachment; filename={}".format(os.path.basename(zip_path))
    return response
@app.route('/segment/destroy_model', methods=['POST'])
def destroy_model_route():
    # 从全局变量中获取模型和其他参数
    global global_state

    # 销毁模型和其他相关变量
    global_state.image = None
    global_state.sam = None
    global_state.predictor = None
    global_state.input_point = None
    global_state.input_label = None
    global_state.input_box_arr = None
    global_state.mask_input = None
    global_state.multimask_output = None
    global_state.input_stop = None
    global_state.image_path = None

    # 清空PyTorch未使用的缓存
    if torch.cuda.is_available():
        torch.cuda.empty_cache()

    return jsonify({"message": "模型销毁成功"})

def allowed_file(filename):
    ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif'}
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS